/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.lvyh.lightframe.tracer.core;

import com.lvyh.lightframe.tracer.core.util.TraceIdGenerator;
import io.opentracing.References;
import io.opentracing.ScopeManager;
import io.opentracing.Span;
import io.opentracing.SpanContext;
import io.opentracing.Tracer;
import io.opentracing.Tracer.SpanBuilder;
import io.opentracing.tag.Tag;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class SimpleSpanBuilder implements Tracer.SpanBuilder {
    private static final Logger logger = LoggerFactory.getLogger(SimpleSpanBuilder.class);

    public static final String ROOT_SPAN_ID = "0";

    private final ScopeManager scopeManager;

    private final String operationName;

    private final Map<String, String> tags;

    private boolean ignoreActiveSpan;

    private SpanContext parentContext;

    private long startTime;

    private SimpleTracer simpleTracer;

    private List<SpanReferenceRelation> references = Collections.emptyList();

    public SimpleSpanBuilder(ScopeManager scopeManager, String operationName) {
        this.scopeManager = scopeManager;
        this.operationName = operationName;
        this.tags = new HashMap<String, String>();
    }

    @Override
    public Tracer.SpanBuilder asChildOf(SpanContext parent) {
        return addReference(References.CHILD_OF, parent);
    }

    @Override
    public Tracer.SpanBuilder asChildOf(Span parentSpan) {
        if (parentSpan == null) {
            return this;
        }
        return addReference(References.CHILD_OF, parentSpan.context());
    }

    public Tracer.SpanBuilder newChild(Span parentSpan) {
        return asChildOf(parentSpan);
    }

    @Override
    public Tracer.SpanBuilder ignoreActiveSpan() {
        ignoreActiveSpan = true;
        return this;
    }

    @Override
    public Tracer.SpanBuilder addReference(String referenceType, SpanContext referencedContext) {
        if (referencedContext == null) {
            return this;
        }
        if (!(referencedContext instanceof SimpleSpanContext)) {
            return this;
        }
        if (!References.CHILD_OF.equals(referenceType)
                && !References.FOLLOWS_FROM.equals(referenceType)) {
            return this;
        }
        if (references.isEmpty()) {
            // Optimization for 99% situations, when there is only one parent
            references = Collections.singletonList(new SpanReferenceRelation(
                    (SimpleSpanContext) referencedContext, referenceType));
        } else {
            if (references.size() == 1) {
                //To ensure order
                references = new ArrayList<>(references);
            }
            references.add(new SpanReferenceRelation(
                    (SimpleSpanContext) referencedContext, referenceType));
        }
        return this;
    }

    @Override
    public Tracer.SpanBuilder withTag(String key, String value) {
        if (key != null && value != null) {
            tags.put(key, value);
        }
        return this;
    }

    @Override
    public Tracer.SpanBuilder withTag(String key, boolean value) {
        if (key != null) {
            tags.put(key, Boolean.toString(value));
        }
        return this;
    }

    @Override
    public Tracer.SpanBuilder withTag(String key, Number value) {
        if (key != null && value != null) {
            tags.put(key, value.toString());
        }
        return this;
    }

    @Override
    public <T> SpanBuilder withTag(Tag<T> key, T value) {
        if (key != null && value != null) {
            tags.put(key.getKey(), value.toString());
        }
        return this;
    }

    @Override
    public Tracer.SpanBuilder withStartTimestamp(long microseconds) {
        startTime = microseconds;
        return this;
    }

    @Override
    public Span start() {
        SimpleSpanContext simpleSpanContext;
        if (this.references != null && this.references.size() > 0) {
            //Parent context exist
            simpleSpanContext = this.createChildContext();
        } else {
            //Start with new root span context
            simpleSpanContext = this.createRootSpanContext();
        }

        for (Map.Entry<String, String> tag : tags.entrySet()) {
            simpleSpanContext.getBaggage().put(tag.getKey(), tag.getValue());
        }

        long begin = this.startTime > 0 ? this.startTime : System.currentTimeMillis();
        SimpleSpan simpleSpan = new SimpleSpan(this.simpleTracer, begin,
                this.references, this.operationName, simpleSpanContext, this.tags);

        logger.info("print traceId:{},spanId:{}", simpleSpan.getSimpleSpanContext().getTraceId(), simpleSpan.getSimpleSpanContext().getSpanId());
        return simpleSpan;
    }


    private SimpleSpanContext createRootSpanContext() {
        String traceId = TraceIdGenerator.generate();
        return new SimpleSpanContext(traceId, ROOT_SPAN_ID, "");
    }

    private SimpleSpanContext createChildContext() {
        SimpleSpanContext preferredReference = preferredReference();

        SimpleSpanContext sofaTracerSpanContext = new SimpleSpanContext(
                preferredReference.getTraceId(), preferredReference.nextChildContextId(),
                preferredReference.getSpanId(), preferredReference.isSampled());
        return sofaTracerSpanContext;
    }

    private SimpleSpanContext preferredReference() {
        SpanReferenceRelation preferredReference = references.get(0);
        for (SpanReferenceRelation reference : references) {
            // childOf takes precedence as a preferred parent
            String referencedType = reference.getReferenceType();
            if (References.CHILD_OF.equals(referencedType)
                    && !References.CHILD_OF.equals(preferredReference.getReferenceType())) {
                preferredReference = reference;
                break;
            }
        }
        return preferredReference.getSimpleSpanContext();
    }

    public Span doStart(Object dispatcher) {
        Span span = new SimpleSpan(dispatcher, baggageItems()).considerStart(startTime).setOperationName(operationName);
        for (Map.Entry<String, String> tag : tags.entrySet()) {
            span.setTag(tag.getKey(), tag.getValue());
        }
        return span;
    }

    private Iterable<Map.Entry<String, String>> baggageItems() {
        if (parentContext != null) { // prefer explicit parent
            return parentContext.baggageItems();
        } else if (!ignoreActiveSpan) {
            Span span = scopeManager.activeSpan();
            if (span != null) {
                return span.context().baggageItems();
            }
        }
        return Collections.<String, String>emptyMap().entrySet();
    }
}